

import os
import re
import sys
import json
import numpy
import threading
import multiprocessing

import cv2

from Module.xRayDetection.xray.logs import Logger
from Module.xRayDetection.xray.util import stringify, get_devices
from Module.xRayDetection.xray.ftp import FTP, put_queue, save_cache, check_folder


def get_region(i, im):
    if i == 0:
        xmin = 0
        xmax = im.shape[1] // 2
        ymin = 0
        ymax = im.shape[0] // 2
    if i == 1:
        xmin = im.shape[1] // 2
        xmax = im.shape[1]
        ymin = 0
        ymax = im.shape[0] // 2
    if i == 2:
        xmin = 0
        xmax = im.shape[1] // 2
        ymin = im.shape[0] // 2
        ymax = im.shape[0]
    if i == 3:
        xmin = im.shape[1] // 2
        xmax = im.shape[1]
        ymin = im.shape[0] // 2
        ymax = im.shape[0]
    return xmin, xmax, ymin, ymax


class OCR:

    def __init__(self, host='', user='', passwd='', ftp=None, logger=None,
        loop=True, sub_queue=None, main_queue=None, 
        cache_path='', save_dir='', remote_dir='', upload_dir='', *args, **kwargs):

        root = os.path.dirname(__file__)

        self.loop = loop

        self.sub_queue = sub_queue
        self.main_queue = main_queue

        self.save_dir = save_dir or os.path.join(root, 'downloads')
        self.upload_dir = upload_dir
        self.remote_dir = remote_dir
        self.cache_dir = os.path.dirname(cache_path)
        self.logger = logger or Logger(os.path.join(self.save_dir, 'logs'))
        self.ftp = ftp or FTP(host=host if not os.path.exists(remote_dir) else '', user=user, passwd=passwd, logger=self.logger)

        cfg_ = stringify(self.__dict__)
        self.logger.log(self.logger.INFO, json.dumps(cfg_, indent=4))

        self.cache = cfg_


    def process_ocr(self, path='', i=0, **kwargs):

        bio = self.ftp.download(download_path=path) if not os.path.exists(path) else open(path, 'rb').read()
        im = numpy.asarray(bytearray(bio.getvalue() if not os.path.exists(path) else bio), dtype='uint8')
        try:
            im = cv2.imdecode(im, cv2.IMREAD_COLOR)
            list_img = []
            for i in range(4):
                xmin, xmax, ymin, ymax = get_region(i, im)
                img = im[ymin:ymax, xmin:xmax, ...]
                list_img.append(img)
            barcode = path.split('/')[-1].split('_')[0]
        except Exception as e:
            print(e)
        
        roi = kwargs['roi']

        for pos, im_ in enumerate(list_img):
            im_roi = im_[roi[1]:roi[3], roi[0]:roi[2], :]
            img_hsv = cv2.cvtColor(im_roi, cv2.COLOR_BGR2HSV)

            if img_hsv[:,:,1:2].mean() > 160 or img_hsv[:,:,1:2].mean() < 9:
                continue

            lower_blue = numpy.array([90,110,110])
            upper_blue = numpy.array([120,255,255])
            mask2 = cv2.inRange(img_hsv, lower_blue, upper_blue)
            # lower mask (0-10)
            lower_red = numpy.array([0,110,110])
            upper_red = numpy.array([10,255,255])
            mask0 = cv2.inRange(img_hsv, lower_red, upper_red)

            # upper mask (170-180)
            lower_red = numpy.array([170,110,110])
            upper_red = numpy.array([180,255,255])
            mask1 = cv2.inRange(img_hsv, lower_red, upper_red)

            lower_green = numpy.array([10,110,110])
            upper_green = numpy.array([80,255,255])
            mask3 = cv2.inRange(img_hsv, lower_green, upper_green)

            # join my masks
            mask = mask0+mask1+mask2+mask3

            try:
                import pytesseract
                string_ = pytesseract.image_to_string(mask, config='--psm 6 --oem 3').replace('\n', ' ')
                string = re.sub('B', '8', string_)
                string = re.sub('E', '6', string)
                string = re.sub('Z', '2', string)
                string = re.sub('\?', '7', string)
                string = re.sub('S|\$', '5', string)
                string = re.sub('d|D|G', '0', string)
                string = re.sub(',\.|,|\.+', '.', string)
                string = re.sub('\d*\+|i|U|L', '1', string)
                min_value = re.search('(\n|^).*?\w* *=* *(\d+\.? *\d+)m*', string)[2]
                max_value = re.search('(\n|^).*?\w* *=* *(\d+\.? *\d+)m*', string[string.index(min_value)+len(min_value):])[2]
                
                min_value = float(min_value.replace(' ', ''))
                max_value = float(max_value.replace(' ', ''))

                min_value = min_value / 100 if min_value > 10 else min_value
                max_value = max_value / 100 if max_value > 10 else max_value

                open(os.path.join(kwargs['dir_'], f"{kwargs['root'].replace('/', '_')}.txt"), 'a').write(f'{barcode},{min_value},{max_value},{pos},{path},{string_}\n')
            except Exception as e:
                cv2.imwrite(os.path.join(kwargs['dir_'], path.split('/')[-1]+f'{pos}.jpg'), im_roi)
                cv2.imwrite(os.path.join(kwargs['dir_'], path.split('/')[-1]+f'hsv{img_hsv[:,:,1:2].mean()}.jpg'), im)
                print(str(e))


    def run(self):

        self.ftp.loop_func(
            self.process_ocr, loop=self.loop, dir_=self.save_dir, items=[self.remote_dir], cache_dir=self.upload_dir,
            main_queue=self.main_queue, sub_queue=self.sub_queue, cache=self.cache)



def run(**kwargs:dict):
    ocr = OCR(**kwargs)
    ocr.run()

def test(args, **kwargs):

    if hasattr(sys, 'gettrace') and sys.gettrace() is not None:
        kwargs['interval'] = 1
        run(**kwargs)
        args = []

    multiprocessing.set_start_method('spawn')

    for arg in args:
        kwargs['pc'] = arg['pc']
        kwargs['device'] = arg['device']
        os.environ['CUDA_VISIBLE_DEVICES'] = str(arg['device'])
        multiprocessing.Process(target=run, kwargs=kwargs).start()
